package fem2.material;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;

import java.util.ArrayList;

import math2.LinAlgUtilities;
import math2.MathUtilities;
import math2.MutableDouble;
import fem2.Element;
import fem2.MeshPart;
import fem2.Variable;
import fem2.element.StructuralElement;
import fem2.enu.State;
import fem2.enu.VariableName;

/**
 * Von Mises (J2 plasticity) material model for: <br>
 * + Plane strain and 3d<br>
 * + Associative flow rule<br>
 * + Associative isotropic hardening<br>
 * + Ziegler kinematic hardening rule<br>
 * + Implicit backward Euler & return mapping
 * 
 * @author hbui
 * 
 */
public class VonMisesMM extends InelasticMaterialModel {

	protected FluencyCriteria isotropicLaw, kinematicLaw;

	public VonMisesMM(double E, double nu, double thickness, double rho, State ss,
			FluencyCriteria isotropicLaw, FluencyCriteria kinematicLaw) {
		super(E, nu, thickness, rho, ss);

		if (getState() == State.PLANE_STRESS) {
			throw new Error(
					"Because of the complexity, this material model does not support plane stress version. Please use VonMisesMMps instead");
		}

		this.isotropicLaw = isotropicLaw;
		this.kinematicLaw = kinematicLaw;
	}

	@Override
	public void Initialize(Element e) {
		super.Initialize(e);

		int dim = e.getMeshPart().getDimension();
		int ntrs = dim * (dim + 1) / 2;
		int ntrs1 = (dim == 2 ? ntrs + 1 : ntrs);

		ArrayList<Variable> vars = getVariables(e);
		for (Variable v : vars) {
			v.registerScalarVariable(VariableName.ALPHA, true);// strain-like
																// variable
			v.registerVectorVariable(VariableName.BETA, ntrs1, true);// stress-like
																		// variable
			v.registerVectorVariable(VariableName.PLASTIC_STRAIN, ntrs1, true);
			/*
			 * save the variable to create history
			 */
			v.save();
		}
	}

	@Override
	public void InitalizeNonlinearIteration(Element e) {
	}

	@Override
	public void FinalizeNonlinearIteration(Element e) {
		StructuralElement se = (StructuralElement) e;
		MeshPart mp = e.getMeshPart();
		int dim = mp.getDimension();
		int ntrs = dim * (dim + 1) / 2;
		int ntrs1 = (dim == 2 ? ntrs + 1 : ntrs);
		int ng = e.getNumIntegrationPoints();
		ArrayVector eedtrial = new ArrayVector(ntrs1);
		ArrayVector strial = new ArrayVector(ntrs1);

		if (dim == 2) {

			/*************************************************************
			 * Von Mises plane strain
			 *************************************************************/

			for (int k = 0; k < ng; k++) {
				double[] xi = e.getIntegrationPoint(k);
				Variable var = getVariables(e).get(k);
				double alphaOld = getVariables(e).get(k).getScalarVariable(VariableName.ALPHA, -1);
				MutableDouble alpha = var.getScalarVariable(VariableName.ALPHA);
				IVector betaOld = var.getVectorVariable(VariableName.BETA, -1);
				IVector beta = var.getVectorVariable(VariableName.BETA);
				IVector epOld = var.getVectorVariable(VariableName.PLASTIC_STRAIN, -1);
				IVector ep = var.getVectorVariable(VariableName.PLASTIC_STRAIN);
				IVector epsilon = var.getVectorVariable(VariableName.STRAIN);
				IMatrix C = getConstitutiveMatrix(e, k);
				IVector sigma = se.getStress(k);
				/*
				 * compute strain
				 */
				BLAM.copy(se.computeStrain(xi), epsilon);
				/*
				 * compute trial deviatoric strain
				 */
				double eevtrial = LinAlgUtilities.trace(epsilon);
				eedtrial.set(0, epsilon.get(0) - eevtrial / 3);
				eedtrial.set(1, epsilon.get(1) - eevtrial / 3);
				eedtrial.set(2, epsilon.get(2));
				eedtrial.set(3, -eevtrial / 3);
				/*
				 * compute trial elastic strain. Note: for Von Mises
				 * tr(plastic_strain) = 0, so epOld = dev(epOld). But it's not
				 * the case for other plasticity model.
				 */
				// IVector eedtrial = (IVector) etrial.clone();
				BLAM.add(-1.0, epOld, eedtrial);
				/*
				 * compute trial deviatoric stress.
				 */
				strial.set(0, 2 * mu * eedtrial.get(0));
				strial.set(1, 2 * mu * eedtrial.get(1));
				strial.set(2, mu * eedtrial.get(2));
				strial.set(3, 2 * mu * eedtrial.get(3));
				/*
				 * compute trial relative stress.
				 */
				IVector xitrial = (IVector) strial.clone();
				BLAM.add(-1.0, betaOld, xitrial);
				/*
				 * check yield condition
				 */
				double normxi = Math.sqrt(xitrial.get(0) * xitrial.get(0) + xitrial.get(1)
						* xitrial.get(1) + 2 * xitrial.get(2) * xitrial.get(2) + xitrial.get(3)
						* xitrial.get(3));
				double oY = isotropicLaw.computeFluencyValue(alphaOld);
				double ftrial = normxi - sqrt_2_3 * oY;
				if (ftrial < 0) {
					/*
					 * update internal variables as trial state
					 */
					alpha.setValue(alphaOld);
					BLAM.copy(betaOld, beta);
					BLAM.copy(epOld, ep);
					/*
					 * update tangent operator
					 */
					BLAM.copy(this.De, C);
					/*
					 * compute trial elastic strain
					 */
					LinAlgUtilities.add(-1.0, ep, epsilon);
					/*
					 * compute stress
					 */
					BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, C, epsilon, 0.0, sigma);

					// if (e.getEchoLevel().contain(EchoLevelType.AUX4)) {
					// Debugger.watch("nonplastic element id = " + e.getId());
					// }

				} else {
					/*
					 * compute N
					 */
					IVector N = (IVector) xitrial.clone();
					BLAM.scale(1 / normxi, N);
					/*
					 * compute delta lambda
					 */
					double dlambda = computeDeltaLambda(normxi, alphaOld);
					/*
					 * update alpha
					 */
					alpha.setValue(alphaOld + sqrt_2_3 * dlambda);
					/*
					 * update beta
					 */
					double H1 = kinematicLaw.computeFluencyValue(alphaOld);
					double H2 = kinematicLaw.computeFluencyValue(alpha.getValue());
					BLAM.copy(betaOld, beta);
					BLAM.add(sqrt_2_3 * (H2 - H1), N, beta);
					/*
					 * update plastic strain
					 */
					epOld = var.getVectorVariable(VariableName.PLASTIC_STRAIN, -1);
					BLAM.copy(epOld, ep);
					BLAM.add(dlambda, N, ep);
					/*
					 * update stress
					 */
					sigma.set(0, kappa * eevtrial + strial.get(0) - 2 * mu * dlambda * N.get(0));
					sigma.set(1, kappa * eevtrial + strial.get(1) - 2 * mu * dlambda * N.get(1));
					sigma.set(2, strial.get(2) - 2 * mu * dlambda * N.get(2));
					/*
					 * update elastoplastic tangent
					 */
					double theta = 1 - 2 * mu * dlambda / normxi;
					double dH = kinematicLaw.computeFluencyDerivative(alpha.getValue());
					double dK = isotropicLaw.computeFluencyDerivative(alpha.getValue());
					double thetabar = 1 / (1 + (dH + dK) / (3 * mu)) - (1 - theta);
					BLAM.zero(C);
					BLAM.add(kappa - 2.0 / 3 * mu * theta, LinAlgUtilities.OneByOne2d, C);
					BLAM.add(2 * mu * theta, LinAlgUtilities.Rank4Identity2d, C);
					LinAlgUtilities.multiply(-2 * mu * thetabar, N, N, C);
				}

			}
		}

		/*************************************************************
		 * Von Mises 3d
		 *************************************************************/

		else if (dim == 3) {
			for (int k = 0; k < ng; k++) {
				double[] xi = e.getIntegrationPoint(k);
				Variable intVar = getVariables(e).get(k);
				double alphaOld = getVariables(e).get(k).getScalarVariable(VariableName.ALPHA, -1);
				MutableDouble alpha = intVar.getScalarVariable(VariableName.ALPHA);
				IVector betaOld = intVar.getVectorVariable(VariableName.BETA, -1);
				IVector beta = intVar.getVectorVariable(VariableName.BETA);
				IVector epOld = intVar.getVectorVariable(VariableName.PLASTIC_STRAIN, -1);
				IVector ep = intVar.getVectorVariable(VariableName.PLASTIC_STRAIN);
				IMatrix C = getConstitutiveMatrix(e, k);
				IVector sigma = se.getStress(k);
				/*
				 * compute strain
				 */
				IVector epsilon = se.computeStrain(xi);
				/*
				 * compute trial deviatoric strain tensor
				 */
				double volStrain = LinAlgUtilities.trace(epsilon);
				eedtrial.set(0, epsilon.get(0) - volStrain / 3);
				eedtrial.set(1, epsilon.get(1) - volStrain / 3);
				eedtrial.set(2, epsilon.get(2) - volStrain / 3);
				eedtrial.set(3, epsilon.get(3));
				eedtrial.set(4, epsilon.get(4));
				eedtrial.set(5, epsilon.get(5));
				/*
				 * compute trial elastic strain
				 */
				// IVector eedtrial = (IVector) eedtrial.clone();
				BLAM.add(-1.0, epOld, eedtrial);
				/*
				 * compute trial deviatoric stress
				 */
				strial.set(0, 2 * mu * eedtrial.get(0));
				strial.set(1, 2 * mu * eedtrial.get(1));
				strial.set(2, 2 * mu * eedtrial.get(2));
				strial.set(3, mu * eedtrial.get(3));
				strial.set(4, mu * eedtrial.get(4));
				strial.set(5, mu * eedtrial.get(5));
				/*
				 * compute trial relative stress
				 */
				IVector xitrial = (IVector) strial.clone();
				BLAM.add(-1.0, betaOld, xitrial);
				double normxi = Math.sqrt(xitrial.get(0) * xitrial.get(0) + xitrial.get(1)
						* xitrial.get(1) + xitrial.get(2) * xitrial.get(2) + 2 * xitrial.get(3)
						* xitrial.get(3) + 2 * xitrial.get(4) * xitrial.get(4) + 2 * xitrial.get(5)
						* xitrial.get(5));
				/*
				 * check yield condition
				 */
				double oY = isotropicLaw.computeFluencyValue(alphaOld);
				double ftrial = normxi - sqrt_2_3 * oY;
				if (ftrial < 0) {
					/*
					 * update internal variables as trial state
					 */
					alpha.setValue(alphaOld);
					BLAM.copy(betaOld, beta);
					BLAM.copy(epOld, ep);
					/*
					 * update tangent operator
					 */
					BLAM.copy(this.De, C);
					/*
					 * compute trial elastic strain
					 */
					BLAM.add(-1.0, ep, epsilon);
					/*
					 * compute stress
					 */
					BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, C, epsilon, 0.0, sigma);
				} else {
					/*
					 * compute N
					 */
					IVector N = (IVector) xitrial.clone();
					BLAM.scale(1 / normxi, N);
					/*
					 * compute deltaGamma
					 */
					double deltaGamma = computeDeltaLambda(normxi, alphaOld);
					/*
					 * update alpha
					 */
					alpha.setValue(alphaOld + sqrt_2_3 * deltaGamma);
					/*
					 * update beta
					 */
					double H1 = kinematicLaw.computeFluencyValue(alphaOld);
					double H2 = kinematicLaw.computeFluencyValue(alpha.getValue());
					BLAM.copy(betaOld, beta);
					BLAM.add(sqrt_2_3 * (H2 - H1), N, beta);
					/*
					 * update plastic strain
					 */
					epOld = intVar.getVectorVariable(VariableName.PLASTIC_STRAIN, -1);
					BLAM.copy(epOld, ep);
					BLAM.add(deltaGamma, N, ep);
					/*
					 * update stress
					 */
					sigma.set(0, kappa * volStrain + strial.get(0) - 2 * mu * deltaGamma * N.get(0));
					sigma.set(1, kappa * volStrain + strial.get(1) - 2 * mu * deltaGamma * N.get(1));
					sigma.set(2, kappa * volStrain + strial.get(2) - 2 * mu * deltaGamma * N.get(2));
					sigma.set(3, strial.get(3) - 2 * mu * deltaGamma * N.get(3));
					sigma.set(4, strial.get(4) - 2 * mu * deltaGamma * N.get(4));
					sigma.set(5, strial.get(5) - 2 * mu * deltaGamma * N.get(5));
					/*
					 * update elastoplastic tangent
					 */
					double theta = 1 - 2 * mu * deltaGamma / normxi;
					double dH = kinematicLaw.computeFluencyDerivative(alpha.getValue());
					double dK = isotropicLaw.computeFluencyDerivative(alpha.getValue());
					double thetabar = 1 / (1 + (dH + dK) / (3 * mu)) - (1 - theta);
					BLAM.zero(C);
					BLAM.add(kappa - 2.0 / 3 * mu * theta, LinAlgUtilities.OneByOne3d, C);
					BLAM.add(2 * mu * theta, LinAlgUtilities.Rank4Identity3d, C);
					LinAlgUtilities.multiply(-2 * mu * thetabar, N, N, C);
				}
			}
		}
	}

	protected double computeDeltaLambda(double normxi, double alpha) {
		double dlambda = 0.0;
		double a = alpha;
		while (true) {
			double H1 = kinematicLaw.computeFluencyValue(alpha);
			double H2 = kinematicLaw.computeFluencyValue(a);
			double g = -sqrt_2_3 * isotropicLaw.computeFluencyValue(a) + normxi
					- (2 * mu * dlambda + sqrt_2_3 * (H2 - H1));
			if (Math.abs(g) < MathUtilities.tol6) {
				break;
			}
			double dH = kinematicLaw.computeFluencyDerivative(a);
			double dK = isotropicLaw.computeFluencyDerivative(a);
			double dg = -2 * mu * (1 + (dH + dK) / (3 * mu));
			dlambda -= g / dg;
			a = alpha + sqrt_2_3 * dlambda;
		}
		return dlambda;
	}

}
